package com.xxx.ratelimiter;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

import com.google.common.util.concurrent.RateLimiter;

/**
 * guava提供的基于令牌桶算法限流器
 * 主要功能：提供一个稳定的速率。
 * 实现方式：通过限制请求流入的速度，计算请求需要等待的合适的时间阈值。RateLimiter通过限制后面请求的等待时间，来支持一定程度的突发请求(预消费)。
 * 工作原理：
 * 		例如：对每秒产生1个令牌的RateLimiter，每当有一个没有使用的令牌，就将storedPermits加1。
 * 		如果RateLimiter在10秒都没有被使用，则storedPermits变成10.0。
 *  	此时，如果一个请求到来，并请求三个令牌(acquire(3))，将从storedPermits中取出3个令牌为其服务，storedPermits变为7.0。
 *  	如果这个请求之后立马又有一个请求到来，并请求10个令牌，将从storedPermits剩余的7个令牌给这个请求，剩下还需要3个令牌，则从RateLimiter新产生的令牌中获取。
 *  	由于RateLimiter每秒只会新产生1个令牌，即上面这个请求还需要的3个请求就要求其等待3秒钟。
 */
public class GuavaRateLimiter {
	
	private String apiKey; // api key
	private double permitsPerSecond; // 每秒处理的最大请求数量
	private Map<String, RateLimiter> limiterMap = new ConcurrentHashMap<>(); // 限流器map

	public String getApiKey() {
		return apiKey;
	}

	public void setApiKey(String apiKey) {
		this.apiKey = apiKey;
	}

	public double getPermitsPerSecond() {
		return permitsPerSecond;
	}

	public void setPermitsPerSecond(double permitsPerSecond) {
		this.permitsPerSecond = permitsPerSecond;
	}
	
	/**
	 * 构造方法
	 * @param apiKey
	 * @param permitsPerSecond
	 */
	public GuavaRateLimiter(String apiKey, double permitsPerSecond) {
		this.apiKey = apiKey;
		this.permitsPerSecond = permitsPerSecond;
		RateLimiter limiter = RateLimiter.create(permitsPerSecond); // 创建guava的限流器
		limiterMap.put(apiKey, limiter); // 将创建好的限流器放入map中
	}
	
	/**
	 * 判断是否能拿到令牌桶中的token
	 * @param ratelimiter
	 * @param apiKey
	 * @return true：拿到了token，说明请求次数没达到限流阈值；false：没有拿到token，说明请求次数达到限流阈值
	 */
	public boolean acquire(String apiKey) {
		RateLimiter limiter = limiterMap.get(apiKey); // 从map中拿到该apiKey对应的限流器
		return limiter.tryAcquire(); // 实际调用的是guava的RateLimiter的tryAcquire(1, 0, MICROSECONDS)方法，非阻塞，没拿到会直接返回；如果要阻塞等待，使用acquire()方法
	}


	/**
	 * 测试guava的RateLimiter的基本用法
	 * @param args
	 * @throws InterruptedException
	 */
	public static void main(String[] args) throws InterruptedException {
//		// case1: 初识RateLimiter
//		RateLimiter limiter = RateLimiter.create(1); // 创建一个每隔1s产生一个token的限流器
//	    for(int n=1; n<20; n=n+2) {
////	        System.out.println("wait before(s)=" + System.currentTimeMillis()/1000 + ", n=" + n + ", wait spent(s)=" + waitTime + ", wait after(s)=" + System.currentTimeMillis()/1000);
//	        System.out.println("n======================"+n);
//	        System.out.println("wait before(s)=" + System.currentTimeMillis()/1000);
//	        double waitTime = limiter.acquire(n); // 同时获取n个token，模拟同时有n个并发请求
//	        System.out.println("wait  spent(s)=" + waitTime);
//	        System.out.println("wait  after(s)=" + System.currentTimeMillis()/1000);
//	    }
//	    // RateLimiter支持预消费，比如在acquire(5)时，等待时间是3秒，原因是上一个获取令牌时预消费了3个令牌，故需要等待3*1秒，然后又预消费了5个令牌，以此类推
	      
		// case2: 
		// 每隔0.5s才会创建一个token放入桶中，如果桶里没有足够的token，则需要等待。
//		double permitsPerSecond = 2.0d;
//		RateLimiter limiter = RateLimiter.create(permitsPerSecond);  // 限流QPS=2，即每隔0.5s左右创建一个token（1个请求对应1个token）
		
		// case2.1
//		// 预先创建一些token（令牌）
//		System.out.println("睡眠1s开始");
//		Thread.sleep(1000);
//		System.out.println("睡眠1s结束");
//		// 本次睡眠结束时剩余token数：2
//		// acquire()，默认为1个请求，返回请被阻塞的时间（单位：秒）
//		System.out.println("currentTime(s)=" + System.currentTimeMillis()/1000 + "\twaitTime(s)=" + limiter.acquire()); //不会等待（使用睡眠时生成的token）
//		System.out.println("currentTime(s)=" + System.currentTimeMillis()/1000 + "\twaitTime(s)=" + limiter.acquire()); //不会等待（使用睡眠时生成的token）
//		System.out.println("currentTime(s)=" + System.currentTimeMillis()/1000 + "\twaitTime(s)=" + limiter.acquire()); //不会等待（预支后一次的token）
//		System.out.println("currentTime(s)=" + System.currentTimeMillis()/1000 + "\twaitTime(s)=" + limiter.acquire(1)); //等待0.5秒，前一次的债务转移
//		System.out.println("currentTime(s)=" + System.currentTimeMillis()/1000 + "\twaitTime(s)=" + limiter.acquire(1)); //等待0.5秒，前一次的债务转移
		
//		// case2.2
//		System.out.println("睡眠2s开始");
//		Thread.sleep(2000);  // 这里无论睡眠多久，最多只保留2个token（maxPermits = maxBurstSeconds * permitsPerSecond）
//		System.out.println("睡眠2s结束");
//		// 本次睡眠结束时剩余token数：2（默认只保留1s时间范围的token）
//		System.out.println("currentTime(s)=" + System.currentTimeMillis()/1000 + "\twaitTime(s)=" + limiter.acquire(1)); //不会等待（使用睡眠时生成的token）
//		System.out.println("currentTime(s)=" + System.currentTimeMillis()/1000 + "\twaitTime(s)=" + limiter.acquire(1)); //不会等待（使用睡眠时生成的token）
//		System.out.println("currentTime(s)=" + System.currentTimeMillis()/1000 + "\twaitTime(s)=" + limiter.acquire(1)); //不会等待（预支后一次的token）
//		System.out.println("currentTime(s)=" + System.currentTimeMillis()/1000 + "\twaitTime(s)=" + limiter.acquire(1)); //等待0.5秒，前一次的债务转移
//		System.out.println("currentTime(s)=" + System.currentTimeMillis()/1000 + "\twaitTime(s)=" + limiter.acquire(1)); //等待0.5秒，前一次的债务转移
//		System.out.println("currentTime(s)=" + System.currentTimeMillis()/1000 + "\twaitTime(s)=" + limiter.acquire(1)); //等待0.5秒，前一次的债务转移
//		System.out.println("currentTime(s)=" + System.currentTimeMillis()/1000 + "\twaitTime(s)=" + limiter.acquire(1)); //等待0.5秒，前一次的债务转移
//		System.out.println("currentTime(s)=" + System.currentTimeMillis()/1000 + "\twaitTime(s)=" + limiter.acquire(1)); //等待0.5秒，前一次的债务转移
//		System.out.println("currentTime(s)=" + System.currentTimeMillis()/1000 + "\twaitTime(s)=" + limiter.acquire(1)); //等待0.5秒，前一次的债务转移
//		System.out.println("currentTime(s)=" + System.currentTimeMillis()/1000 + "\twaitTime(s)=" + limiter.acquire(1)); //等待0.5秒，前一次的债务转移
		
//        // case2.3: 非阻塞方式获取token：如果没有拿到，直接返回
//        Thread.sleep(1000); // 睡眠了1秒钟，此时桶里肯定会有token，如果去掉睡眠，则桶里不一定会有token
//        System.out.println(System.currentTimeMillis()/1000 + "\t" + limiter.tryAcquire()); // true
//        System.out.println(System.currentTimeMillis()/1000 + "\t" + limiter.tryAcquire()); // true
        
        // case4: 阻塞方式获取token：等待一段时间，如果还是没有拿到，则返回
//		System.out.println("currentTime(s)=" + System.currentTimeMillis()/1000 + "\t" + limiter.tryAcquire(1100, TimeUnit.MILLISECONDS)); // 尝试等待1100ms true
//		System.out.println("currentTime(s)=" + System.currentTimeMillis()/1000 + "\t" + limiter.tryAcquire(400, TimeUnit.MILLISECONDS));  // 尝试等待400ms false
	}
}
